package com.kinyx.framework.security;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.interfaces.RSAKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

import com.kinyx.framework.Configs;
import com.kinyx.framework.annotation.assist.CanNotNull;
import com.kinyx.framework.annotation.assist.CanNull;
import com.kinyx.framework.annotation.assist.ReturnNotNull;
import com.kinyx.framework.annotation.assist.Tester;
import com.kinyx.framework.annotation.assist.Throws;
import com.kinyx.framework.enums.impl.ErrorCodeX;
import com.kinyx.framework.exception.KRuntimeException;
import com.kinyx.framework.file.KFileReader;

import lombok.AccessLevel;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.ToString;

/**
 * RSA编码工具类
 */
@Tester("com.kinyx.framework.security.KRSA_Tester_")
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public final class KRSA {

	@ReturnNotNull
	//	@Throws(at = "RSA.rsa()", value = "NoSuchAlgorithmException", when = "KeyFactory.getInstance(String)")
	private static KeyFactory rsa() {
		try {
			return KeyFactory.getInstance("RSA");
		} catch (final NoSuchAlgorithmException e) {// never happens
			throw new KRuntimeException(e, ErrorCodeX.NeverHappens);
		}
	}

	@ReturnNotNull
	private static int calcMaxLength(@CanNotNull final boolean encrypt, @CanNotNull final RSAKey key) {
		final int keySize = key.getModulus().toString(2).length() / 8;
		return encrypt ? keySize - 11 : keySize;
	}

	@ReturnNotNull
	//	@Throws(at = "RSA.cipher(int, Key)", value = "NoSuchAlgorithmException", when = "Cipher.getInstance(String)")
	//	@Throws(at = "RSA.cipher(int, Key)", value = "NoSuchPaddingException", when = "Cipher.getInstance(String)")
	@Throws(at = "RSA.cipher(int, Key)", value = "InvalidKeyException", when = "Cipher.init(int, Key)")
	private static Cipher cipher(@CanNotNull final int opmode, @CanNotNull final Key key) {
		Cipher cipher;
		try {
			cipher = Cipher.getInstance("RSA");
		} catch (final NoSuchAlgorithmException e) {// never happens
			throw new KRuntimeException(e, ErrorCodeX.NeverHappens);
		} catch (final NoSuchPaddingException e) {// never happens
			throw new KRuntimeException(e, ErrorCodeX.NeverHappens);
		}
		try {
			cipher.init(opmode, key);
		} catch (final InvalidKeyException e) {
			throw new KRuntimeException(e);
		}
		return cipher;
	}

	@ReturnNotNull
	@Throws(at = "RSA.doFinal(Cipher, byte[], int)", value = "IllegalBlockSizeException", when = "Cipher.doFinal(byte[])")
	@Throws(at = "RSA.doFinal(Cipher, byte[], int)", value = "BadPaddingException", when = "Cipher.doFinal(byte[])")
	private static byte[] doFinal(@CanNotNull final Cipher cipher, @CanNotNull final byte[] bytes, @CanNotNull final int maxCipherLength) {
		try {
			if (bytes.length > maxCipherLength) {// 长度超过最大可解析长度，进行分段处理。
				ByteArrayOutputStream baos = new ByteArrayOutputStream();
				for (int idx = 0; idx < bytes.length; idx += maxCipherLength) {
					byte[] cache;
					cache = cipher.doFinal(bytes, idx, Math.min(maxCipherLength, bytes.length - idx));
					baos.write(cache, 0, cache.length);
				}
				final byte[] byteArray = baos.toByteArray();
				try {
					baos.close();
				} catch (final IOException e) {// ignore this
				} finally {
					if (baos != null) {
						baos = null;
					}
				}
				return byteArray;
			}
			return cipher.doFinal(bytes);
		} catch (final IllegalBlockSizeException e) {
			throw new KRuntimeException(e);
		} catch (final BadPaddingException e) {
			throw new KRuntimeException(e);
		}
	}

	/**
	 * 加密
	 * @param plaintext 明文
	 * @param publicKey 公钥
	 * @param charset 字符集
	 * @return 密文
	 */
	//	@Throws(at = "RSA.rsa()", value = "NoSuchAlgorithmException", when = "KeyFactory.getInstance(String)")
	//	@Throws(at = "RSA.cipher(int, Key)", value = "NoSuchAlgorithmException", when = "Cipher.getInstance(String)")
	//	@Throws(at = "RSA.cipher(int, Key)", value = "NoSuchPaddingException", when = "Cipher.getInstance(String)")
	@Throws(at = "RSA.cipher(int, Key)", value = "InvalidKeyException", when = "Cipher.init(int, Key)")
	@Throws(at = "RSA.doFinal(Cipher, byte[], int)", value = "IllegalBlockSizeException", when = "Cipher.doFinal(byte[])")
	@Throws(at = "RSA.doFinal(Cipher, byte[], int)", value = "BadPaddingException", when = "Cipher.doFinal(byte[])")
	@Throws(at = "RSA.encrypt(String, String, String)", value = "InvalidKeySpecException", when = "KeyFactory.generatePublic(KeySpec)")
	@Throws(at = "RSA.encrypt(String, String, String)", value = "UnsupportedEncodingException", when = "String.getBytes(String)")
	@ReturnNotNull
	public static String encrypt(@CanNotNull final String plaintext, @CanNotNull final String publicKey, @CanNull(defaultValue = "Configs.default_RSA_charset") final String charset) {
		if (charset == null) { return encrypt(plaintext, publicKey, Configs.default_RSA_charset); }
		RSAPublicKey key;
		try {
			key = (RSAPublicKey) rsa().generatePublic(new X509EncodedKeySpec(KBASE64.decode(publicKey)));
		} catch (final InvalidKeySpecException e) {
			throw new KRuntimeException(e);
		}
		try {
			return KBASE64.encodeToString(doFinal(cipher(Cipher.ENCRYPT_MODE, key), plaintext.getBytes(charset), calcMaxLength(true, key)));
		} catch (final UnsupportedEncodingException e) {
			throw new KRuntimeException(e);
		}
	}

	/**
	 * 解密
	 * @param ciphertext 密文
	 * @param privateKey 私钥
	 * @param charset 字符集
	 * @return 明文
	 */
	//	@Throws(at = "RSA.rsa()", value = "NoSuchAlgorithmException", when = "KeyFactory.getInstance(String)")
	//	@Throws(at = "RSA.cipher(int, Key)", value = "NoSuchAlgorithmException", when = "Cipher.getInstance(String)")
	//	@Throws(at = "RSA.cipher(int, Key)", value = "NoSuchPaddingException", when = "Cipher.getInstance(String)")
	@Throws(at = "RSA.cipher(int, Key)", value = "InvalidKeyException", when = "Cipher.init(int, Key)")
	@Throws(at = "RSA.doFinal(Cipher, byte[], int)", value = "IllegalBlockSizeException", when = "Cipher.doFinal(byte[])")
	@Throws(at = "RSA.doFinal(Cipher, byte[], int)", value = "BadPaddingException", when = "Cipher.doFinal(byte[])")
	@Throws(at = "RSA.decrypt(String, String, String)", value = "InvalidKeySpecException", when = "KeyFactory.generatePrivate(KeySpec)")
	@Throws(at = "RSA.decrypt(String, String, String)", value = "UnsupportedEncodingException", when = "new String(String)")
	@ReturnNotNull
	public static String decrypt(@CanNotNull final String ciphertext, @CanNotNull final String privateKey, @CanNull(defaultValue = "Configs.default_RSA_charset") final String charset) {
		if (charset == null) { return decrypt(ciphertext, privateKey, Configs.default_RSA_charset); }
		RSAPrivateKey key;
		try {
			key = (RSAPrivateKey) rsa().generatePrivate(new PKCS8EncodedKeySpec(KBASE64.decode(privateKey)));
		} catch (final InvalidKeySpecException e) {
			throw new KRuntimeException(e);
		}
		try {
			return new String(doFinal(cipher(Cipher.DECRYPT_MODE, key), KBASE64.decode(ciphertext), calcMaxLength(false, key)), charset);
		} catch (final UnsupportedEncodingException e) {
			throw new KRuntimeException(e);
		}
	}

	/**
	 * 秘钥对
	 */
	@Getter
	@AllArgsConstructor
	@ToString
	public static class KeyPair {

		private final String privateKey;
		private final String publicKey;

	}

	/**
	 * 创建秘钥对
	 * @param keysize 秘钥长度
	 * @param key 辅助加密
	 * @param charset 字符集
	 * @return 秘钥对
	 */
	@Throws(at = "KBase64.encode(byte[], String)", value = "UnsupportedEncodingException", when = "new String(byte[], String)")
	//	@Throws(at = "RSA.generateKeyPair(int, String, String)", value = "NoSuchAlgorithmException", when = "KeyPairGenerator.getInstance(String)")
	@Throws(at = "RSA.generateKeyPair(int, String, String)", value = "UnsupportedEncodingException", when = "String.getBytes(String)")
	@ReturnNotNull
	public static KeyPair generateKeyPair(@CanNotNull final int keysize, @CanNull final String key, @CanNull(defaultValue = "Configs.default_RSA_charset") final String charset) {
		if (charset == null) { return generateKeyPair(keysize, key, Configs.default_RSA_charset); }
		KeyPairGenerator kpg;
		try {
			kpg = KeyPairGenerator.getInstance("RSA");
		} catch (final NoSuchAlgorithmException e) {// never happens
			throw new KRuntimeException(e, ErrorCodeX.NeverHappens);
		}
		try {
			kpg.initialize(keysize, key == null ? new SecureRandom() : new SecureRandom(key.getBytes(charset)));
		} catch (final UnsupportedEncodingException e) {
			throw new KRuntimeException(e);
		}
		final java.security.KeyPair kp = kpg.generateKeyPair();
		return new KeyPair(KBASE64.encode(kp.getPrivate().getEncoded(), charset), KBASE64.encode(kp.getPublic().getEncoded(), charset));
	}

	private static final String BEGIN_PRIVATE_KEY = "-----BEGIN PRIVATE KEY-----";
	private static final String END_PRIVATE_KEY = "-----END PRIVATE KEY-----";
	private static final String BEGIN_PUBLIC_KEY = "-----BEGIN PUBLIC KEY-----";
	private static final String END_PUBLIC_KEY = "-----END PUBLIC KEY-----";

	@Throws(at = "KReader.read(InputStream, String)", value = "UnsupportedEncodingException", when = "new InputStreamReader(InputStream, String)")
	@Throws(at = "KReader.read(InputStream, String)", value = "IOException", when = "InputStreamReader.read()")
	@Throws(at = "KFileUtils.checkFile(File)", value = "ErrorCode0.FileDoesNotExists", when = "!File.exists()")
	@Throws(at = "KFileUtils.checkFile(File)", value = "ErrorCode0.IsNotAFile", when = "!File.isFile()")
	@Throws(at = "KFileUtils.checkReadable(File)", value = "ErrorCode0.FileIsNotReadable", when = "!File.canRead()")
	//	@Throws(at = "KFileReader.newFileInputStream(File)", value = "FileNotFoundException", when = "new FileInputStream(File)")
	@ReturnNotNull
	private static String readKey(final boolean pk, @CanNotNull final String fileName, @CanNull(defaultValue = "Configs.default_KReader_charset") final String charset, final boolean retainComments, final boolean retainLineBreaker) {
		final String[] lines = KFileReader.read(new File(fileName), charset).split("\r\n");
		final StringBuilder sb = new StringBuilder();
		final String lineFirst = lines[0];
		if ((pk && lineFirst.equals(BEGIN_PRIVATE_KEY)) || (!pk && lineFirst.equals(BEGIN_PUBLIC_KEY))) {
			if (retainComments) {
				sb.append(lineFirst);
				if (retainLineBreaker) { sb.append("\n"); }
			}
		} else {
			sb.append(lineFirst);
			if (retainLineBreaker) { sb.append("\n"); }
		}
		for (int i = 1; i < (lines.length - 1); i++) {
			sb.append(lines[i]);
			if (retainLineBreaker) { sb.append("\n"); }
		}
		final String lineLast = lines[lines.length - 1];
		if ((pk && lineLast.equals(END_PRIVATE_KEY)) || (!pk && lineLast.equals(END_PUBLIC_KEY))) {
			if (retainComments) {
				sb.append(lineLast);
			}
		} else {
			sb.append(lineLast);
		}
		if (sb.charAt(sb.length() - 1) == '\n') {
			sb.deleteCharAt(sb.length() - 1);
		}
		return sb.toString();
	}

	/**
	 * 从文件中读取RSA私钥
	 * @param fileName 文件名（全路径）
	 * @param charset 字符集
	 * @param retainComments true: 保留注释; false: 去除注释;
	 * @param retainLineBreaker true: 保留换行符; false: 去除换行符;
	 * @return RSA私钥
	 */
	@Throws(at = "KReader.read(InputStream, String)", value = "UnsupportedEncodingException", when = "new InputStreamReader(InputStream, String)")
	@Throws(at = "KReader.read(InputStream, String)", value = "IOException", when = "InputStreamReader.read()")
	@Throws(at = "KFileUtils.checkFile(File)", value = "ErrorCode0.FileDoesNotExists", when = "!File.exists()")
	@Throws(at = "KFileUtils.checkFile(File)", value = "ErrorCode0.IsNotAFile", when = "!File.isFile()")
	@Throws(at = "KFileUtils.checkReadable(File)", value = "ErrorCode0.FileIsNotReadable", when = "!File.canRead()")
	//	@Throws(at = "KFileReader.newFileInputStream(File)", value = "FileNotFoundException", when = "new FileInputStream(File)")
	@ReturnNotNull
	public static String readPrimaryKey(@CanNotNull final String fileName, @CanNull(defaultValue = "Configs.default_KReader_charset") final String charset, final boolean retainComments, final boolean retainLineBreaker) {
		return readKey(true, fileName, charset, retainComments, retainLineBreaker);
	}

	/**
	 * 从文件中读取RSA公钥
	 * @param fileName 文件名（全路径）
	 * @param charset 字符集
	 * @param retainComments true: 保留注释; false: 去除注释;
	 * @param retainLineBreaker true: 保留换行符; false: 去除换行符;
	 * @return RSA公钥
	 */
	@Throws(at = "KReader.read(InputStream, String)", value = "UnsupportedEncodingException", when = "new InputStreamReader(InputStream, String)")
	@Throws(at = "KReader.read(InputStream, String)", value = "IOException", when = "InputStreamReader.read()")
	@Throws(at = "KFileUtils.checkFile(File)", value = "ErrorCode0.FileDoesNotExists", when = "!File.exists()")
	@Throws(at = "KFileUtils.checkFile(File)", value = "ErrorCode0.IsNotAFile", when = "!File.isFile()")
	@Throws(at = "KFileUtils.checkReadable(File)", value = "ErrorCode0.FileIsNotReadable", when = "!File.canRead()")
	//	@Throws(at = "KFileReader.newFileInputStream(File)", value = "FileNotFoundException", when = "new FileInputStream(File)")
	@ReturnNotNull
	public static String readPublicKey(@CanNotNull final String fileName, @CanNull(defaultValue = "Configs.default_KReader_charset") final String charset, final boolean retainComments, final boolean retainLineBreaker) {
		return readKey(false, fileName, charset, retainComments, retainLineBreaker);
	}

}
